//
//  VideoEncoder.swift
//  replaykit
//
//  Created by zhang on 2025/7/31.
//

import Foundation
import VideoToolbox
import CoreMedia
import CoreVideo

class VideoEncoder: NSObject {
    
    // MARK: - Properties
    private var session: VTCompressionSession?
    private var isEncoding = false
    
    // 编码配置
    private let width: Int
    private let height: Int
    private let fps: Int
    private let bitrate: Int
    
    // 回调
    var onEncodedData: ((Data, UInt32) -> Void)?
    var onError: ((String) -> Void)?
    
    // MARK: - Initialization
    init(width: Int = 1280, height: Int = 720, fps: Int = 30, bitrate: Int = 2000000) {
        self.width = width
        self.height = height
        self.fps = fps
        self.bitrate = bitrate
        super.init()
    }
    
    deinit {
        stopEncoder()
    }
    
    // MARK: - Encoder Setup
    func setupEncoder() -> Bool {
        print("VideoEncoder: 开始设置H.264编码器")
        
        // 创建编码会话
        let status = VTCompressionSessionCreate(
            allocator: kCFAllocatorDefault,
            width: Int32(width),
            height: Int32(height),
            codecType: kCMVideoCodecType_H264,
            encoderSpecification: nil,
            imageBufferAttributes: nil,
            compressedDataAllocator: nil,
            outputCallback: outputCallback,
            refcon: UnsafeMutableRawPointer(Unmanaged.passUnretained(self).toOpaque()),
            compressionSessionOut: &session
        )
        
        guard status == noErr, let session = session else {
            print("VideoEncoder: 创建编码会话失败: \(status)")
            onError?("创建编码会话失败")
            return false
        }
        
        // 配置编码参数
        configureEncoder(session: session)
        
        // 开始编码会话
        let startStatus = VTCompressionSessionPrepareToEncodeFrames(session)
        guard startStatus == noErr else {
            print("VideoEncoder: 准备编码失败: \(startStatus)")
            onError?("准备编码失败")
            return false
        }
        
        isEncoding = true
        print("VideoEncoder: H.264编码器设置成功")
        return true
    }
    
    private func configureEncoder(session: VTCompressionSession) {
        // 设置实时编码
        VTSessionSetProperty(session, key: kVTCompressionPropertyKey_RealTime, value: kCFBooleanTrue)
        
        // 设置比特率
        let bitrateValue = bitrate as CFNumber
        VTSessionSetProperty(session, key: kVTCompressionPropertyKey_AverageBitRate, value: bitrateValue)
        
        // 设置最大比特率
        let maxBitrateValue = (bitrate * 2) as CFNumber
        VTSessionSetProperty(session, key: kVTCompressionPropertyKey_DataRateLimits, value: maxBitrateValue)
        
        // 设置帧率
        let fpsValue = fps as CFNumber
        VTSessionSetProperty(session, key: kVTCompressionPropertyKey_ExpectedFrameRate, value: fpsValue)
        
        // 设置关键帧间隔
        let keyframeInterval = fps * 2 as CFNumber // 每2秒一个关键帧
        VTSessionSetProperty(session, key: kVTCompressionPropertyKey_MaxKeyFrameInterval, value: keyframeInterval)
        
        // 设置编码配置文件 (Baseline Profile)
        VTSessionSetProperty(session, key: kVTCompressionPropertyKey_ProfileLevel, value: kVTProfileLevel_H264_Baseline_AutoLevel)
        
        print("VideoEncoder: 编码参数配置完成")
    }
    
    // MARK: - Encoding
    func encodeFrame(_ pixelBuffer: CVPixelBuffer, timestamp: CMTime) {
        guard let session = session, isEncoding else {
            print("VideoEncoder: 编码器未准备就绪")
            return
        }
        
        let status = VTCompressionSessionEncodeFrame(
            session,
            imageBuffer: pixelBuffer,
            presentationTimeStamp: timestamp,
            duration: CMTime.invalid,
            frameProperties: nil,
            sourceFrameRefcon: nil,
            infoFlagsOut: nil
        )
        
        if status != noErr {
            print("VideoEncoder: 编码帧失败: \(status)")
            onError?("编码帧失败")
        }
    }
    
    func stopEncoder() {
        guard let session = session else { return }
        
        print("VideoEncoder: 停止编码器")
        VTCompressionSessionCompleteFrames(session, untilPresentationTimeStamp: CMTime.invalid)
        VTCompressionSessionInvalidate(session)
        self.session = nil
        isEncoding = false
    }
    
    // MARK: - Output Callback
    private let outputCallback: VTCompressionOutputCallback = { (outputCallbackRefCon, sourceFrameRefCon, status, flags, sampleBuffer) in
        guard let refCon = outputCallbackRefCon else { return }
        let encoder = Unmanaged<VideoEncoder>.fromOpaque(refCon).takeUnretainedValue()
        encoder.handleEncodedFrame(status: status, flags: flags, sampleBuffer: sampleBuffer)
    }
    
    private func handleEncodedFrame(status: OSStatus, flags: VTEncodeInfoFlags, sampleBuffer: CMSampleBuffer?) {
        guard status == noErr, let sampleBuffer = sampleBuffer else {
            print("VideoEncoder: 编码回调错误: \(status)")
            return
        }
        
        // 检查是否为关键帧
        let attachments = CMSampleBufferGetSampleAttachmentsArray(sampleBuffer, createIfNecessary: false)
        var isKeyFrame = false
        if let attachments = attachments as? [[CFString: Any]], let attachment = attachments.first {
            if let dependsOnOthers = attachment[kCMSampleAttachmentKey_NotSync] as? Bool {
                isKeyFrame = !dependsOnOthers
            }
        }
        
        // 获取编码后的数据
        guard let dataBuffer = CMSampleBufferGetDataBuffer(sampleBuffer) else {
            print("VideoEncoder: 无法获取编码数据")
            return
        }
        
        // 获取数据长度
        var dataLength = 0
        var dataPointer: UnsafeMutablePointer<Int8>?
        let status = CMBlockBufferGetDataPointer(dataBuffer, atOffset: 0, lengthAtOffsetOut: nil, totalLengthOut: &dataLength, dataPointerOut: &dataPointer)
        
        guard status == kCMBlockBufferNoErr, let dataPointer = dataPointer else {
            print("VideoEncoder: 无法获取数据指针")
            return
        }
        
        // 创建数据
        let data = Data(bytes: dataPointer, count: dataLength)
        
        // 获取时间戳
        let presentationTime = CMSampleBufferGetPresentationTimeStamp(sampleBuffer)
        let timestamp = UInt32(CMTimeGetSeconds(presentationTime) * 1000) // 转换为毫秒
        
        print("VideoEncoder: 编码成功 - 大小: \(data.count) bytes, 时间戳: \(timestamp), 关键帧: \(isKeyFrame)")
        
        // 回调编码后的数据
        onEncodedData?(data, timestamp)
    }
    
    // MARK: - Status
    var isEncodingStatus: Bool {
        return isEncoding
    }
} 